package com.wmzdq.aiservice.controller;

import com.alibaba.cloud.ai.dashscope.chat.DashScopeChatModel;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.client.advisor.api.Advisor;
import org.springframework.ai.document.Document;
import org.springframework.ai.rag.Query;
import org.springframework.ai.rag.advisor.RetrievalAugmentationAdvisor;
import org.springframework.ai.rag.generation.augmentation.ContextualQueryAugmenter;
import org.springframework.ai.rag.retrieval.search.DocumentRetriever;
import org.springframework.ai.rag.retrieval.search.VectorStoreDocumentRetriever;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Qualifier;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;

import java.util.List;

@RestController
@RequestMapping("/rag")
@Slf4j
public class RagController {

    @Autowired
    private VectorStore vectorStore;

    @Qualifier("dashscopeChatModel")
    DashScopeChatModel chatModel;
//    @Autowired
//    ChatClient.Builder chatClientBuilder;

    @GetMapping("/test")
    public String test(@RequestParam(value = "query", defaultValue = "你好，很高兴认识你，能简单介绍一下自己吗？")String query){
        DocumentRetriever retriever = VectorStoreDocumentRetriever.builder()
                .vectorStore(vectorStore)
                .similarityThreshold(0.5)    // 设置相似度阈值
                .topK(3)                     // 返回前3个最相关的文档
                //.filterExpression(filterExpression.build())
                .build();

        ChatClient chatClient = ChatClient.builder(chatModel)
                .defaultSystem("你是一个专业的电商售后客服，请根据客户需要进行回答")
                .build();

        Query q = new Query(query);
        List<Document> list= retriever.apply(q);

        Advisor advisor = RetrievalAugmentationAdvisor.builder()
                .queryAugmenter(ContextualQueryAugmenter.builder()
                        .allowEmptyContext(true)
                        .build())
                .documentRetriever(retriever)
                .build();

        String response = chatClient.prompt()
                .user(query)
                .advisors(advisor)
                .call()
                .content();
        return response;
    }
}
